#!/bin/bash
# PRISM 高性能训练 - 快速启动脚本

echo "=================================================="
echo "    PRISM 高性能训练 - 快速启动"
echo "=================================================="
echo ""

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 检查Python环境
echo -e "${BLUE}[1/5] 检查Python环境...${NC}"
if ! command -v python &> /dev/null; then
    echo -e "${RED}错误: 未找到Python${NC}"
    exit 1
fi
echo -e "${GREEN}✓ Python已安装${NC}"

# 检查必要的包
echo -e "${BLUE}[2/5] 检查依赖包...${NC}"
python -c "import torch, ultralytics, optuna" 2>/dev/null
if [ $? -ne 0 ]; then
    echo -e "${YELLOW}警告: 部分依赖包未安装${NC}"
    echo "运行: pip install -r requirements.txt"
else
    echo -e "${GREEN}✓ 依赖包已安装${NC}"
fi

# 检查数据集
echo -e "${BLUE}[3/5] 检查数据集...${NC}"
if [ -d "dataset" ]; then
    echo -e "${GREEN}✓ 数据集目录存在${NC}"
else
    echo -e "${RED}错误: 未找到dataset目录${NC}"
    exit 1
fi

# 检查GPU
echo -e "${BLUE}[4/5] 检查GPU...${NC}"
python -c "import torch; print(f'GPU可用: {torch.cuda.is_available()}'); print(f'GPU数量: {torch.cuda.device_count()}')"
if [ $? -eq 0 ]; then
    echo -e "${GREEN}✓ GPU检查完成${NC}"
else
    echo -e "${YELLOW}警告: GPU检查失败，将使用CPU训练（非常慢）${NC}"
fi

# 创建必要的目录
echo -e "${BLUE}[5/5] 创建输出目录...${NC}"
mkdir -p high_performance_results
mkdir -p optimization_results
mkdir -p weights
echo -e "${GREEN}✓ 目录创建完成${NC}"

echo ""
echo "=================================================="
echo "    环境检查完成，选择训练模式:"
echo "=================================================="
echo ""
echo "  1) 完整训练流程（推荐，12-24小时）"
echo "  2) 仅Stage 1训练（2-4小时）"
echo "  3) 仅Stage 2训练（3-6小时）"
echo "  4) 仅超参数优化（6-12小时）"
echo "  5) 仅评估（10分钟）"
echo "  6) 快速测试（所有阶段，少量epoch）"
echo "  0) 退出"
echo ""
read -p "请选择 [0-6]: " choice

case $choice in
    1)
        echo -e "${GREEN}启动完整训练流程...${NC}"
        python train_high_performance.py --mode full
        ;;
    2)
        echo -e "${GREEN}启动Stage 1训练...${NC}"
        python train_high_performance.py --mode stage1
        ;;
    3)
        echo -e "${GREEN}启动Stage 2训练...${NC}"
        python train_high_performance.py --mode stage2
        ;;
    4)
        echo -e "${GREEN}启动超参数优化...${NC}"
        read -p "输入试验次数（建议30-50）: " n_trials
        python train_high_performance.py --mode optimize --n-trials ${n_trials:-30}
        ;;
    5)
        echo -e "${GREEN}启动评估...${NC}"
        python train_high_performance.py --mode evaluate
        ;;
    6)
        echo -e "${YELLOW}快速测试模式（减少epoch）...${NC}"
        # 创建临时配置
        cat > config_quick_test.yaml << EOF
stage1:
  epochs: 10
  batch_size: 16
  img_size: 640
  patience: 5

stage2:
  epochs: 10
  batch_size: 32
  learning_rate: 0.0001
  patience: 5
  model:
    use_fpn: false
    use_roi_align: false

optimization:
  n_trials: 5
  timeout: null
EOF
        python train_high_performance.py --mode full --config config_quick_test.yaml
        ;;
    0)
        echo "退出"
        exit 0
        ;;
    *)
        echo -e "${RED}无效选择${NC}"
        exit 1
        ;;
esac

echo ""
echo "=================================================="
echo "    训练完成！"
echo "=================================================="
echo ""
echo "查看结果:"
echo "  - 训练历史: high_performance_results/training_history.json"
echo "  - 最终评估: high_performance_results/final_evaluation/"
echo "  - 优化结果: optimization_results/"
echo "  - 日志文件: training_*.log"
echo ""
